import sqlite3 as sql
from AFTdata import AFTdata


db_dir = "~/Caucasus/db"



def read_aft(samples=None,pull_detrital=True):
    """Reads aft data into AFTdata class. Samples is a list of sample names, or a comma delimited striing of sample names. If none, all samples are read"""
    if isinstance(samples,str): samples = samples.split(',')
    aft=[]
    table = []
    conn = sql.connect('%s/data.sqlite' % db_dir)
    query = "SELECT sample,report,report_id FROM ages WHERE method='AFT'"
    if not pull_detrital: query+=" AND comment IS NOT 'detrital'"
    if samples is None:
        table = conn.cursor().execute(query).fetchall()
    else:
        query+=" AND sample=?"  # safety first
        for sample in samples:
            table.append(conn.cursor().execute(query, (sample,)).fetchone())
    conn.close()

    for sample,report,report_id in table:
        aft.append(AFTdata("%s/AFT/A%s%sA.AGE" % (db_dir,report,report_id),
                           "%s/AFT/L%s%sA.A2G" % (db_dir,report,report_id), 
                      new_zeta = int(report)>=1145))  # I think Ray is back to the standard zeta
        aft[-1].name = sample
    if len(aft)==1: return aft[0]

    return aft
